#script will process samples with three color virus and two color fiducials on the same image
#it will select reference that gives smallest alignemnt among references
#inputs: red, green and blue images. Plus a 16 by 16 image of anything.
#this script outputs reference, which is a list of the distances obtained from teh cross
#correlations of the fiducials using the best reference. Reference is RG, virus is RGB
#070612
import numpy, pylab
from scipy.ndimage import gaussian_filter
from scipy.signal import resample, fftconvolve
import math


######open_files#####
def read_data(file_name, z, x, y):
    return numpy.fromfile(file_name, dtype = numpy.uint16
                   ).reshape(z,x,y).astype(numpy.float)


data_green = read_data('g_1.raw', 81, 256, 256)
data_red = read_data('r_1.raw', 81, 256, 256)
data_blue = read_data('b_1.raw', 81, 256, 256)
green_reference = read_data('g_1.raw', 81, 256, 256)
red_reference = read_data('r_1.raw', 81, 256, 256)
reference_shape = read_data('r_2ref.raw', 81, 16, 16)

d = reference_shape.shape
upsample_factor = (10, 6, 6)


##pylab.close('all')
##fig = pylab.figure()
##pylab.imshow(red_reference.sum(axis=0), cmap=pylab.cm.gray, interpolation='nearest')
##fig.show()

         
############################################
######FIND_PEAK_POSITIONS###################
############################################

    ################blue########################   
#psf_radius = 8
psf_radius = max(reference_shape.shape[1:])//2
print "psf radius", psf_radius

def find_peaks(data_name, small_gauss_factor = 0.1, large_gauss_factor = 10, std_dev = 5, std_dev_sat = 3):
    data_summed = data_name.sum(axis=0)
    data_filtered = gaussian_filter(data_summed, small_gauss_factor) - gaussian_filter(data_summed, large_gauss_factor)
    data_peaks = data_filtered*(data_filtered > ((std_dev*data_filtered.std()+data_filtered.mean())))
                                 #and (data_filtered < (std_dev_sat*data_filtered.std())+data_filtered.mean())))
    #data_peaks = data_filtered*(data_filtered > (std_dev*data_filtered.std())+data_filtered.mean())
    data_peaks_copy = data_peaks.copy()
    data_peak_locations = []
    while data_peaks_copy.max() > 0:
        x, y = numpy.unravel_index(numpy.argmax(data_peaks_copy), data_peaks_copy.shape)
        data_peak_locations.append((x, y))
        data_peaks_copy[max(x-psf_radius, 0):x+psf_radius,
                        max(y-psf_radius, 0):y+psf_radius] = 0
    #pylab.close('all')
    #fig = pylab.figure()
    #pylab.imshow(data_peaks, cmap=pylab.cm.gray, interpolation='nearest')
    #fig.show()
    return data_peak_locations


blue_peak_locations = find_peaks(data_blue, 0.01, 10, 0.5, 3)
green_peak_locations = find_peaks(data_green, 0.1, 10, 5, 3)
red_peak_locations = find_peaks(data_red, 0.1, 10, 5, 3)
redREF_peak_locations = find_peaks(red_reference, 0.1, 10, 5, 3)
greenREF_peak_locations = find_peaks(green_reference, 0.1, 10, 5, 3)

#################################################
###############PEAK_SORTING######################
#################################################

#print "blue peak locations", blue_peak_locations
#print "green peak locations", green_peak_locations
#print "red peak locations", red_peak_locations

    #################red_and_green################

#print "green peak locations", green_peak_locations
#print "red peak locations", red_peak_locations

    ############peaks_in_red_and_green_and_blue####

def peak_label_sorting(tolerance, data1_peak_locations, data2_peak_locations):
    peaks_in_data1_data2 = []
    for x, y in data1_peak_locations:
        for x_b, y_b in data2_peak_locations:
            #Look for a b-match
            if (abs(x - x_b) <= tolerance and
                abs(y - y_b) <= tolerance):
                #We found a b-match
                peaks_in_data1_data2.append((x, y))
                break #Stop searching for a b-match
    return peaks_in_data1_data2

def blue_subtraction(tolerance, rg_peaks, rgb_peaks):
    rg_only = []
    for x, y in rg_peaks:
        isMatch = False
        for x_b, y_b in rgb_peaks:
            if (abs(x - x_b) <= tolerance and
                abs(y - y_b) <= tolerance):
                isMatch = True
                break
        if not isMatch:
            rg_only.append((x, y))
    return rg_only
    

peaks_in_red_green = peak_label_sorting(2, red_peak_locations, green_peak_locations)
peaks_in_rgb = peak_label_sorting(2, peaks_in_red_green, blue_peak_locations)
peaks_in_rg_only = blue_subtraction(2, peaks_in_red_green, peaks_in_rgb)

#print "peaks in red green", len(peaks_in_red_green)
print "peaks in rg_only", len(peaks_in_rg_only)
print "peaks in RGB", len(peaks_in_rgb)


###############################################
#######UPSAMPLING##############################
###############################################

def resample_3(array, n, window='hamming'):
    s = array.shape
    array = array.copy()
    for i in range(3):
        array = resample(array, n[i]*s[i], t=None, axis=i, window=window)
    return array


##################REFERENCE_SQUARES_FFT_AND_UPSAMPLE#########################


def gaussian_filter_fft(data_name, small_gauss_factor = 0.3, large_gauss_factor = 10):
    filtered_data = gaussian_filter(data_name, small_gauss_factor) - gaussian_filter(data_name, large_gauss_factor)
    return numpy.fft.fftn(filtered_data)

def up_and_cross_data(data_name, data_reference, data_reference_upsampled_fft, peak_coordinates, color):
    data_wrapped_shift_pixels = []
    #print color , "shape", d
    for p in peak_coordinates:
        data_subimage = data_name[:,
                        p[0]-(d[1]//2):p[0]-(d[1]//2)+d[1],
                        p[1]-(d[2]//2):p[1]-(d[2]//2)+d[2]]
        
        if data_subimage.shape != d:
            #print "Skipping"
            continue

        data_subimage_upsampled = resample_3(data_subimage, upsample_factor, window='hamming')

        print "Cross-correlating green..."
        data_cross_correlation = numpy.fft.ifftn(
            data_reference_upsampled_fft *
            numpy.conjugate(numpy.fft.fftn(data_subimage_upsampled)))
        #print "Finding max..."
        data_shift = numpy.unravel_index(data_cross_correlation.argmax(), data_cross_correlation.shape)
        data_wrapped_shift = list(data_shift)
        pixel_size = 50.0
        data_wrapped_shift_pixels_zxy = []
        for i in range(3):
            if data_shift[i] > data_cross_correlation.shape[i]//2:
                data_wrapped_shift[i] -= data_cross_correlation.shape[i]
            data_wrapped_shift_pixels_zxy.append(data_wrapped_shift[i]/float(upsample_factor[i])*pixel_size)
            pixel_size = 134.0
        data_wrapped_shift_pixels.append(data_wrapped_shift_pixels_zxy)
        #print color, "wrapped shift:", data_wrapped_shift

    return data_wrapped_shift_pixels

def distance3D_calc(x1, x2, y1, y2, z1, z2):
    return math.sqrt((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)

def distance2D_calc(x1, x2, y1, y2):
    return math.sqrt((x1 - x2)**2 + (y1 - y2)**2)

    ###########FIND BEST REFERENCE##################
index_ref_distances3D = []
total_index_ref_distances3D = []
index_ref_distances2D = []
total_index_ref_distances2D = []
loop_count = 0
for a in peaks_in_rg_only:
    loop_count = loop_count + 1
    reference_squares = []
    reference_subimage_green = green_reference[:,
                a[0]-(d[1]//2):a[0]-(d[1]//2)+d[1],
                a[1]-(d[2]//2):a[1]-(d[2]//2)+d[2]]
    reference_subimage_red = red_reference[:,
                a[0]-(d[1]//2):a[0]-(d[1]//2)+d[1],
                a[1]-(d[2]//2):a[1]-(d[2]//2)+d[2]]
    if reference_subimage_green.shape != d or reference_subimage_red.shape != d:
        #print "Skipping"
        continue    
    reference_subimage_green_upsampled = resample_3(reference_subimage_green, upsample_factor, window='hamming')
    reference_subimage_green_ftt_upsampled = gaussian_filter_fft(reference_subimage_green_upsampled, 0.3, 10)
    reference_subimage_red_upsampled = resample_3(reference_subimage_red, upsample_factor, window='hamming')
    reference_subimage_red_ftt_upsampled = gaussian_filter_fft(reference_subimage_red_upsampled, 0.3, 10)
    
    green_wrapped_shift_pixels_ref = up_and_cross_data(data_green, reference_subimage_green, reference_subimage_green_ftt_upsampled, peaks_in_rg_only, 'green')
    red_wrapped_shift_pixels_ref = up_and_cross_data(data_red, reference_subimage_red, reference_subimage_red_ftt_upsampled, peaks_in_rg_only, 'red')
    ref_distances2D = []
    ref_distances3D = []
    for i in range(len(red_wrapped_shift_pixels_ref)):
        z1, x1, y1 = red_wrapped_shift_pixels_ref[i]
        z2, x2, y2 = green_wrapped_shift_pixels_ref[i]
        ref_distances2D.append(distance2D_calc(x1, x2, y1, y2))
        ref_distances3D.append(distance3D_calc(x1, x2, y1, y2, z1, z2))
        total_index_ref_distances2D.append(distance2D_calc(x1, x2, y1, y2))
        total_index_ref_distances3D.append(distance3D_calc(x1, x2, y1, y2, z1, z2))
    #print 'ref_distances3d', ref_distances3D
    index_ref_distances2D.append([(loop_count-1),(numpy.average(ref_distances2D))])
    index_ref_distances3D.append([(loop_count-1),(numpy.average(ref_distances3D))])
    #print 'distance calc', distance3D_calc(x1, x2, y1, y2, z1, z2)
#print 'ref_dist', index_ref_distances3D

indexmin = index_ref_distances3D[0][0]
valuemin = index_ref_distances3D[0][1]
for a in index_ref_distances3D:
    if a[1] < valuemin:
        indexmin = a[0]
        valuemin = a[1]

#indexmin = min(index_ref_distances3D[:][1])
#print 'index min', indexmin


reference_subimage_green = green_reference[:,
        peaks_in_rg_only[indexmin][0]-(d[1]//2):peaks_in_rg_only[indexmin][0]-(d[1]//2)+d[1],
        peaks_in_rg_only[indexmin][1]-(d[2]//2):peaks_in_rg_only[indexmin][1]-(d[2]//2)+d[2]]
reference_subimage_red = red_reference[:,
        peaks_in_rg_only[indexmin][0]-(d[1]//2):peaks_in_rg_only[indexmin][0]-(d[1]//2)+d[1],
        peaks_in_rg_only[indexmin][1]-(d[2]//2):peaks_in_rg_only[indexmin][1]-(d[2]//2)+d[2]]

##pylab.close('all')
##fig = pylab.figure()
##pylab.imshow(reference_subimage_red.sum(axis=0), cmap=pylab.cm.gray, interpolation='nearest')
##fig.show()
       
reference_subimage_green_upsampled = resample_3(reference_subimage_green, upsample_factor, window='hamming')
reference_subimage_green_ftt_upsampled = gaussian_filter_fft(reference_subimage_green_upsampled, 0.2, 10)

reference_subimage_red_upsampled = resample_3(reference_subimage_red, upsample_factor, window='hamming')
reference_subimage_red_ftt_upsampled = gaussian_filter_fft(reference_subimage_red_upsampled, 0.2, 10)

###########DISTANCES OF FIDUCIALS USING BEST REFERENCE###########

best_fid_red_shift_pixels = up_and_cross_data(data_green, reference_subimage_green, reference_subimage_green_ftt_upsampled, peaks_in_rg_only, 'green')
best_fid_green_shift_pixels = up_and_cross_data(data_red, reference_subimage_red, reference_subimage_red_ftt_upsampled, peaks_in_rg_only, 'red')
fid_best_distances3D = []
fid_best_distances2D = []

for i in range(len(best_fid_red_shift_pixels)):
    z1, x1, y1 = best_fid_red_shift_pixels[i]
    z2, x2, y2 = best_fid_green_shift_pixels[i]
    fid_best_distances3D.append(distance3D_calc(x1, x2, y1, y2, z1, z2))
    fid_best_distances2D.append(distance2D_calc(x1, x2, y1, y2))


##############UPSAMPLE AND CROSSCORRELATE DATA################
green_wrapped_shift_pixels = up_and_cross_data(data_green, reference_subimage_green, reference_subimage_green_ftt_upsampled, peaks_in_rgb, 'green')
red_wrapped_shift_pixels = up_and_cross_data(data_red, reference_subimage_red, reference_subimage_red_ftt_upsampled, peaks_in_rgb, 'red')
distances3D = []
distances2D = []

for i in range(len(red_wrapped_shift_pixels)):
    z1, x1, y1 = red_wrapped_shift_pixels[i]
    z2, x2, y2 = green_wrapped_shift_pixels[i]
    distances3D.append(distance3D_calc(x1, x2, y1, y2, z1, z2))
    distances2D.append(distance2D_calc(x1, x2, y1, y2))

##############VIEW RESULTS ON SCREEN#########################

print '/nindex and reference distances for 3d', index_ref_distances3D
print '/nindex and reference distances for 2d', index_ref_distances2D
print
for i in distances3D:
    print '%0.2f'%(i),
print

yy = [y for y in distances3D if y < 200]
print len(yy)
average3D = numpy.mean(yy)
print 'average 3D', average3D

print
for i in distances2D:
    print '%0.2f'%(i),
print

xx = [x for x in distances2D if x < 200]
print len(xx)
average2D = numpy.mean(xx)
print 'average 2D', average2D

    ####ref###

zz = [z for z in total_index_ref_distances3D if z < 200]
print len(zz)
ave_ref_3D = numpy.mean(zz)
print 'average reference 3D', ave_ref_3D

oo = [o for o in total_index_ref_distances2D if o < 200]
print len(oo)
ave_ref_2D = numpy.mean(oo)
print 'average reference 2D', ave_ref_2D

    ####best ref####

qq = [q for q in fid_best_distances3D if q < 200]
print len(qq)
ave_best_ref_3D = numpy.mean(qq)
print 'average best reference 3D', ave_best_ref_3D

rr = [r for r in fid_best_distances2D if r < 200]
print len(rr)
ave_best_ref_2D = numpy.mean(rr)
print 'average best reference 2D', ave_best_ref_2D

######SAVE FILES########
file_out3D = open('results3D_r.txt', 'w')
for item in distances3D:
  file_out3D.write("%0.2f\n" % item)
file_out3D.close()

file_out2D = open('results2D_r.txt', 'w')
for item in distances2D:
  file_out2D.write("%0.2f\n" % item)
file_out2D.close()

file_outref3D = open('resultsREF3D.txt', 'w')
for item in total_index_ref_distances3D:
  file_outref3D.write("%0.2f\n" % item)
file_outref3D.close()

file_outref2D = open('resultsREF2D.txt', 'w')
for item in total_index_ref_distances2D:
  file_outref2D.write("%0.2f\n" % item)
file_outref2D.close()

file_outbestref3D = open('resultsBESTREF3D.txt', 'w')
for item in fid_best_distances3D:
  file_outbestref3D.write("%0.2f\n" % item)
file_outbestref3D.close()

file_outbestref2D = open('resultsBESTREF2D.txt', 'w')
for item in fid_best_distances2D:
  file_outbestref2D.write("%0.2f\n" % item)
file_outbestref2D.close()




